package ai.mspbots.poc.es.service;

import ai.mspbots.poc.es.entity.Document;
import ai.mspbots.poc.help.VectorHelper;
import cn.hutool.json.JSONUtil;
import org.elasticsearch.action.search.SearchRequest;
import org.elasticsearch.action.search.SearchResponse;
import org.elasticsearch.client.RequestOptions;
import org.elasticsearch.client.RestHighLevelClient;
import org.elasticsearch.common.lucene.search.function.FunctionScoreQuery;
import org.elasticsearch.index.query.BoolQueryBuilder;
import org.elasticsearch.index.query.MatchQueryBuilder;
import org.elasticsearch.index.query.QueryBuilders;
import org.elasticsearch.index.query.functionscore.FunctionScoreQueryBuilder;
import org.elasticsearch.index.query.functionscore.ScriptScoreFunctionBuilder;
import org.elasticsearch.script.Script;
import org.elasticsearch.script.ScriptType;
import org.elasticsearch.search.SearchHit;
import org.elasticsearch.search.builder.SearchSourceBuilder;
import org.nd4j.linalg.api.ndarray.INDArray;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

import java.io.IOException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

@Service
public class VectorService {

    private static final Logger logger = LoggerFactory.getLogger(VectorService.class);

    @Autowired
    private RestHighLevelClient client;

    public List<Document> search(String text) throws IOException {
        SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();

        // 创建布尔查询
        BoolQueryBuilder boolQuery = new BoolQueryBuilder();

        // 添加模糊匹配查询
        MatchQueryBuilder matchQuery = new MatchQueryBuilder("name", text);
        matchQuery.fuzziness("AUTO");
        boolQuery.should(matchQuery);

        // 准备脚本参数
        Map<String, Object> params = new HashMap<>();
        INDArray indArray = VectorHelper.vectorizedText(text);
        double[] doubleVector = indArray.toDoubleVector();
        logger.info("text = {}, Vector = {}", text, doubleVector);
        params.put("userVectorName", doubleVector);
        // params.put("userVectorDescription", VectorHelper.vectorizedText(text).toDoubleVector());

        // 使用 script_score 查询
        Script script = new Script(
                ScriptType.INLINE,
                "painless",
                "double sum = 0;\n" +
                        "          List userVector = params.userVectorName;\n" +
                        "\n" +
                        "          for (int i = 0; i < userVector.size(); i++) {\n" +
                        "            if (doc['name_vector'].size() == 0) {\n" +
                        "              return 0;\n" +
                        "            }\n" +
                        "            double diff = doc['name_vector'].vectorValue[i] - (double)userVector.get(i);\n" +
                        "            sum += Math.pow(diff, 2);\n" +
                        "          }\n" +
                        "          \n" +
                        "          double distance = Math.sqrt(sum);\n" +
                        "          return Math.max(0, -distance + 10);",
                params
        );
        //
        // Script script = new Script(ScriptType.INLINE, "painless",
        //         "cosineSimilarity(params.userVectorName, 'description_vector') + 1.0", params);

        // searchSourceBuilder.query(QueryBuilders.scriptScoreQuery(QueryBuilders.matchAllQuery(), script));

        ScriptScoreFunctionBuilder scriptScoreFunction = new ScriptScoreFunctionBuilder(script);

        // 将脚本评分查询添加到布尔查询中
        FunctionScoreQueryBuilder functionScoreQuery = new FunctionScoreQueryBuilder(
                boolQuery,
                new FunctionScoreQueryBuilder.FilterFunctionBuilder[]{
                        new FunctionScoreQueryBuilder.FilterFunctionBuilder(scriptScoreFunction)
                }
        ).scoreMode(FunctionScoreQuery.ScoreMode.SUM); // 或者选择其他模式
        searchSourceBuilder.query(QueryBuilders.scriptScoreQuery(QueryBuilders.matchAllQuery(), script));
        // searchSourceBuilder.query(functionScoreQuery);

        // 执行查询
        SearchRequest searchRequest = new SearchRequest("sys_report_vector");
        searchRequest.source(searchSourceBuilder);
        SearchResponse searchResponse = client.search(searchRequest, RequestOptions.DEFAULT);

        List<Document> documents = new ArrayList<>();

        for (SearchHit hit : searchResponse.getHits().getHits()) {
            String source = hit.getSourceAsString();
            Document doc = JSONUtil.toBean(source, Document.class);
            // 继续提取其他字段
            documents.add(doc);
        }
        return documents;
    }
}